#ifndef SRC_MPI_PARTITION_3D
#define SRC_MPI_PARTITION_3D

/**
MIT License

Copyright (c) 2020 Hua Huang

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

#include "log_helper.hpp"

#include <algorithm>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vector>

struct Solution {
  long long pm, pn, pk;
  double value;

  Solution(long long pm, long long pn, long long pk, long long value)
      : pm(pm), pn(pn), pk(pk), value(value) {}
};

static Solution findOptimal(long long p, long long bvm, long long bvn,
                            long long bvk, long long R, bool verbose = false) {
  std::vector<Solution> solutions;
  double a1 = static_cast<double>(R), a2 = bvm * bvk;
  double b1 = static_cast<double>(R), b2 = bvn * bvk;
  double c = bvm * bvn;

  for (long long pm = 1; pm <= p; ++pm) {
    if (p % pm != 0 || pm > bvm) {
      continue;
    }
    long long rest1 = p / pm;
    for (long long pn = 1; pn <= rest1; ++pn) {
      if (rest1 % pn != 0 || pn > bvn) {
        continue;
      }
      long long pk = rest1 / pn;
      if (pk > bvk) {
        continue;
      }

      double val = (a1 * std::min(pn - 1, 1ll) + 1) * a2 * pn +
                   (b1 * std::min(pm - 1, 1ll) + 1) * b2 * pm + c * pk;
      solutions.emplace_back(pm, pn, pk, val);
      if (verbose) {
        printf("pm=%lld pn=%lld pk=%lld target=%f\n", pm, pn, pk, val);
      }
    }
  }

  auto cmp = [](const Solution &a, const Solution &b) {
    return a.value < b.value;
  };

  auto min_sol = min_element(solutions.begin(), solutions.end(), cmp);
  return *min_sol;
}

static double max_min_ratio(const int n_elem, double *a) {
  double maxval = a[0];
  double minval = a[0];
  for (int i = 1; i < n_elem; i++) {
    maxval = (maxval < a[i]) ? a[i] : maxval;
    minval = (minval > a[i]) ? a[i] : minval;
  }
  return (maxval / minval);
}

static void fixed_pair_scale_grid(const int p, const int idx0, const int idx1,
                                  double *size_grid, double *perfect_grid,
                                  int *proc_grid) {
  int idx2 = 0 + 1 + 2 - idx0 - idx1;
  double fixed_ratio = perfect_grid[idx1] / perfect_grid[idx0];
  double p_double = p;
  double perfect_sg[3];
  double perfect_vol = 1.0;
  for (int i = 0; i < 3; i++) {
    perfect_sg[i] = size_grid[i] / perfect_grid[i];
    perfect_vol *= perfect_sg[i];
  }

  // Find a conservative partition as the baseline
  proc_grid[idx0] = (int)floor(perfect_grid[idx0]);
  proc_grid[idx1] = (int)floor(fixed_ratio) * proc_grid[idx0];
  proc_grid[idx2] =
      (int)(floor(p_double / (double)(proc_grid[idx0] * proc_grid[idx1])));
  double sg_proc[3], min_vol = 1.0, surf0, max_surfsum;
  for (int i = 0; i < 3; i++) {
    sg_proc[i] = size_grid[i] / (double)proc_grid[i];
    min_vol *= sg_proc[i];
  }
  surf0 = sg_proc[0] * sg_proc[1] + sg_proc[0] * sg_proc[2] +
          sg_proc[1] * sg_proc[2];
  max_surfsum = surf0 * (double)proc_grid[0] * (double)proc_grid[1] *
                (double)proc_grid[2];

  int proc_grid1[3], valid;
  double sg_proc1[3], vol1, surf1, surfsum1;
  int pg_idx0_lower = (int)ceil(perfect_grid[idx0] * 0.3);
  int pg_idx0_upper = (int)floor(perfect_grid[idx0] * 3.3);
  int ratio_lower = (int)ceil(fixed_ratio * 0.3);
  int ratio_upper = (int)floor(fixed_ratio * 3.3);
  for (int pg_idx0 = pg_idx0_lower; pg_idx0 <= pg_idx0_upper; pg_idx0++) {
    proc_grid1[idx0] = pg_idx0;
    for (int ratio = ratio_lower; ratio <= ratio_upper; ratio++) {
      proc_grid1[idx1] = ratio * proc_grid1[idx0];
      proc_grid1[idx2] =
          (int)floor(p_double / (double)(proc_grid1[idx0] * proc_grid1[idx1]));
      vol1 = 1.0;
      surfsum1 = 1.0;
      for (int i = 0; i < 3; i++) {
        sg_proc1[i] = size_grid[i] / (double)proc_grid1[i];
        vol1 *= sg_proc1[i];
        surfsum1 *= (double)proc_grid1[i];
      }
      surf1 = sg_proc1[0] * sg_proc1[1] + sg_proc1[0] * sg_proc1[2] +
              sg_proc1[1] * sg_proc1[2];
      surfsum1 = surfsum1 * surf1;
      if ((proc_grid1[0]) < 1 || (proc_grid1[1] < 1) || (proc_grid1[2] < 1))
        continue;
      if (vol1 / perfect_vol > 1.25)
        continue;
      if (max_min_ratio(3, sg_proc1) >= 4.0)
        continue;
      if (vol1 < min_vol) {
        min_vol = vol1;
        max_surfsum = surfsum1;
        proc_grid[0] = proc_grid1[0];
        proc_grid[1] = proc_grid1[1];
        proc_grid[2] = proc_grid1[2];
      }
      if ((vol1 < min_vol * 1.05) && (surfsum1 < max_surfsum)) {
        min_vol = vol1;
        max_surfsum = surfsum1;
        proc_grid[0] = proc_grid1[0];
        proc_grid[1] = proc_grid1[1];
        proc_grid[2] = proc_grid1[2];
      }
    }
  }
}

// Calculate the near-optimal 3D decomposition of cuboid of size m * n * k for p
// processes s.t. max(mp, np) is a multiple of min(mp, np) Input parameters:
//   m, n, k : Size of the cuboid
//   p       : Number of processes
// Output parameters:
//   mp, np, kp : Active process grid size
//   rp         : Number of idle processes
static void calc_3d_decomposition(const int p, const int m, const int n,
                                  const int k, int *mp, int *np, int *kp,
                                  int *rp) {
  double size_grid[3], perfect_grid[3];
  int idx_grid[3], proc_grid[3] = {1, 1, 1};
  if ((m <= n) && (n <= k)) {
    size_grid[0] = m;
    size_grid[1] = n;
    size_grid[2] = k;
    idx_grid[0] = 0;
    idx_grid[1] = 1;
    idx_grid[2] = 2;
  }
  if ((m <= k) && (k <= n)) {
    size_grid[0] = m;
    size_grid[1] = k;
    size_grid[2] = n;
    idx_grid[0] = 0;
    idx_grid[1] = 2;
    idx_grid[2] = 1;
  }
  if ((n <= m) && (m <= k)) {
    size_grid[0] = n;
    size_grid[1] = m;
    size_grid[2] = k;
    idx_grid[0] = 1;
    idx_grid[1] = 0;
    idx_grid[2] = 2;
  }
  if ((n <= k) && (k <= m)) {
    size_grid[0] = n;
    size_grid[1] = k;
    size_grid[2] = m;
    idx_grid[0] = 2;
    idx_grid[1] = 0;
    idx_grid[2] = 1;
  }
  if ((k <= m) && (m <= n)) {
    size_grid[0] = k;
    size_grid[1] = m;
    size_grid[2] = n;
    idx_grid[0] = 1;
    idx_grid[1] = 2;
    idx_grid[2] = 0;
  }
  if ((k <= n) && (n <= m)) {
    size_grid[0] = k;
    size_grid[1] = n;
    size_grid[2] = m;
    idx_grid[0] = 2;
    idx_grid[1] = 1;
    idx_grid[2] = 0;
  }

  double d0 = size_grid[0], d1 = size_grid[1], d2 = size_grid[2];

  // One large dimension cases
  double p_double = p;
  double d2d1_ratio = d2 / d1;
  if (d2d1_ratio > p_double)
    proc_grid[2] = p;

  // Two large dimension cases
  double large2_threshold = d1 * d2 / (d0 * d0);
  double k_double = k;
  if ((d2d1_ratio <= p_double) && (p_double <= large2_threshold)) {
    double max_pg_prod = 1.0;
    int ratio_lower = (int)ceil(d2d1_ratio * 0.5);
    int ratio_upper = (int)floor(d2d1_ratio * 2);
    for (int ratio = ratio_lower; ratio <= ratio_upper; ratio++) {
      double ratio_double = ratio;
      int d1p = (int)floor(sqrt(p_double / ratio_double));
      //   int d2p =
      //       (k_double == d0) ? (d1p * ratio) : (int)floor(p_double /
      //       (double)d1p);
      //?! remove the requirement: 'mod(max(mp,np),min(mp,np))=0'
      int d2p = (int)floor(p_double / (double)d1p);
      double d1p_d2p = (double)d1p * d2p;
      if (d1p_d2p > max_pg_prod) {
        max_pg_prod = d1p_d2p;
        proc_grid[1] = d1p;
        proc_grid[2] = d2p;
      }
    }
  }

  // Three large dimension cases
  if (p_double > large2_threshold) {
    double d2d0_ratio = d2 / d0;
    double d1d0_ratio = d1 / d0;
    double d0_perfect = cbrt(p_double / (d2d0_ratio * d1d0_ratio));
    double d1_perfect = d0_perfect * d1d0_ratio;
    double d2_perfect = d0_perfect * d2d0_ratio;
    perfect_grid[0] = d0_perfect;
    perfect_grid[1] = d1_perfect;
    perfect_grid[2] = d2_perfect;
    if (idx_grid[2] == 0)
      fixed_pair_scale_grid(p, 1, 2, &size_grid[0], &perfect_grid[0],
                            &proc_grid[0]);
    if (idx_grid[2] == 1)
      fixed_pair_scale_grid(p, 0, 2, &size_grid[0], &perfect_grid[0],
                            &proc_grid[0]);
    if (idx_grid[2] == 2)
      fixed_pair_scale_grid(p, 0, 1, &size_grid[0], &perfect_grid[0],
                            &proc_grid[0]);
  }

  *mp = proc_grid[idx_grid[0]];
  *np = proc_grid[idx_grid[1]];
  *kp = proc_grid[idx_grid[2]];
  *rp = p - (*mp) * (*np) * (*kp);
  if (*rp * 10 > p) {
    WARNING_PRINTF(
        "Failed to find a good solution, current: %d * %d * %d + %d\n", *mp,
        *np, *kp, *rp);
  }
}

// Calculate the near-optimal 3D decomposition of cuboid of size n * n * k for p
// processes Input parameters:
//   n, k : Size of the cuboid
//   p    : Number of processes
// Output parameters:
//   np, kp : Active process grid size (np * np * kp)
//   rp     : Number of idle processes
static void calc_3d_decomposition_nk(const int p, const int n, const int k,
                                     int *np, int *kp, int *rp) {
  int tmp_mp = 1, tmp_np = 1, tmp_kp = 1, tmp_rp;
  double d = (double)k / (double)n;
  tmp_mp = (int)ceil(cbrt((double)p / d));
  tmp_np = tmp_mp;
  tmp_kp = p / (tmp_mp * tmp_np);
  tmp_rp = p - tmp_mp * tmp_np * tmp_kp;
  *np = tmp_np;
  *kp = tmp_kp;
  *rp = tmp_rp;
  while (tmp_mp > 1) {
    tmp_mp--;
    tmp_np--;
    tmp_kp = p / (tmp_mp * tmp_np);
    tmp_rp = p - tmp_mp * tmp_np * tmp_kp;
    double d1 = (double)tmp_kp / (double)tmp_np;
    if (d1 >= 4.0 * d)
      break;
    if (tmp_rp < *rp) {
      *np = tmp_np;
      *kp = tmp_kp;
      *rp = tmp_rp;
    }
  }
  if (*rp * 10 > p) {
    WARNING_PRINTF(
        "Failed to find a good solution, current: %d * %d * %d + %d\n", *np,
        *np, *kp, *rp);
  }
}

#endif /* SRC_MPI_PARTITION_3D */
